#include "DWMRICore.h"

#include <vector>

using namespace std;

char *fiber_pathname = NULL;
char *dst_pathname = NULL;

void PrintUsage()
{
	printf("Usage: MappingFibers -fiber configure_file -dst dst_file\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		}else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL)
		return -2;

	return 0;
}

void AddFiber(vector<CPoint3I> &map, const int x, const int y, const int z)
{
	for (int i = 0; i < map.size(); ++i) {
		if (map.at(i).x == x && map.at(i).y == y && map.at(i).z == z) {
			return;
		}
	}
	map.push_back(CPoint3I(x, y, z));
}

void MappingFibers(CFiber *fibers, const int w, const int h, const int d, 
	const int z, const int scale, vector<vector<CPoint3I>> &fiberMap)
{
	for (int y = 0; y < h; ++y) { 
		for (int x = 0; x < w; ++x) {
			int i = y * w + x;
			/* mapping every fiber */
			CFiber *f = &(fibers[i]);

			int lastID = -1;
			for (uint k = 1; k < f->m_fCount; ++k) {
				CPoint3F sPt = f->m_pF[k-1];
				CPoint3F ePt = f->m_pF[k];
				int sID = (int(sPt.m_z * scale + 0.5f) * h + int(sPt.m_y * scale + 0.5f)) * w + int(sPt.m_x * scale + 0.5f);
				int eID = (int(ePt.m_z * scale + 0.5f) * h + int(ePt.m_y * scale + 0.5f)) * w + int(ePt.m_x * scale + 0.5f);
				if (sID == eID && sID != lastID) {
					AddFiber(fiberMap[sID], x, y, z);
					lastID = sID;
				} else {
					CVector3F dir;
					dir = ePt - sPt;
					float len = dir.Length();
					dir.Normalize();
					
					for (float l = 0.0; l < len; l += 0.01f) {
						CPoint3F pt = sPt + dir * l;
						int id = (int(pt.m_z * scale + 0.5f) * h + int(pt.m_y * scale + 0.5f)) * w + int(pt.m_x * scale + 0.5f);
						if (id != lastID) {
							AddFiber(fiberMap[id], x, y, z);
							lastID = id;
						}
					}
				}
			}

			for (uint k = 1; k < f->m_bCount; ++k) {
				CPoint3F sPt = f->m_pB[k-1];
				CPoint3F ePt = f->m_pB[k];
				int sID = (int(sPt.m_z * scale + 0.5f) * h + int(sPt.m_y * scale + 0.5f)) * w + int(sPt.m_x * scale + 0.5f);
				int eID = (int(ePt.m_z * scale + 0.5f) * h + int(ePt.m_y * scale + 0.5f)) * w + int(ePt.m_x * scale + 0.5f);
				if (sID == eID && sID != lastID) {
					AddFiber(fiberMap[sID], x, y, z);
					lastID = sID;
				} else {
					CVector3F dir;
					dir = ePt - sPt;
					float len = dir.Length();
					dir.Normalize();
					
					for (float l = 0.0; l < len; l += 0.01f) {
						CPoint3F pt = sPt + dir * l;
						int id = (int(pt.m_z * scale + 0.5f) * h + int(pt.m_y * scale + 0.5f)) * w + int(pt.m_x * scale + 0.5f);
						if (id != lastID) {
							AddFiber(fiberMap[id], x, y, z);
							lastID = id;
						}
					}
				}
			}
		}
	}
}

void GetOneVoxelFibers(char **confPathname, char **vertPathname, vector<CPoint3I> &map)
{
	int size = int(map.size());
	if (size == 0)
		return;

	CFiber *fiber = new CFiber[size];
	for (int i = 0; i < size; ++i) {
		CPoint3I pos = map.at(i);
		CFiber *fibersBuffer = NULL;
		int cw, ch;
		ReadFibers(confPathname[pos.z], vertPathname[pos.z], &fibersBuffer, cw, ch);
		fiber[i].CopyFiber(&(fibersBuffer[pos.y*cw+pos.x]));
		delete[] fibersBuffer;
	}

	SaveFibersAsVTK("C:\\Ziang\\Dataset\\MITK\\temp\\fibers.vtk", fiber, size);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	/* read the configure file */
	char **confPathname = NULL;
	char **vertPathname = NULL;
	int w, h, d, scale;
	float stepSize, maxLength;
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);

	vector<vector<CPoint3I>> fiberMap;
	fiberMap.resize(w*h*d);

	/* mapping fibers for each slice */
	for (int z = 0; z < d; ++z) {
		printf("Mapping fiber %s & %s ... ", confPathname[z], vertPathname[z]);

		CFiber *fibers;
		int cw, ch;
		ReadFibers(confPathname[z], vertPathname[z], &fibers, cw, ch);
		
		MappingFibers(fibers, w, h, d, z, scale, fiberMap);
		
		delete[] fibers;

		printf("Done.\n");
	}

	/* for debug */
	int voxelID = (49 * h + 128) * w + 101;
	GetOneVoxelFibers(confPathname, vertPathname, fiberMap.at(voxelID));

	for (int z = 0; z < d; ++z) {
		delete[] confPathname[z];
		delete[] vertPathname[z];
	}
	delete[] confPathname;
	delete[] vertPathname;

	float *count = new float[w*h*d];
	for (int i = 0; i < w*h*d; ++i) {
		count[i] = float(fiberMap.at(i).size());
	}
	CScalar mappingCount;
	mappingCount.CreateScalar(w, h, d, count);
	mappingCount.SaveScalarFile(dst_pathname);



	delete[] count;

	return 0;
}
